//
//  TreeNode.swift
//  数据结构与算法
//
//  Created by ZERO on 2021/1/9.
//

import Foundation

/// 前缀树
class TreeNode {
    /// 节点的名称，在前缀树里是单个字母
    var label: Character
    /// 使用哈希映射存放子结点。哈希便于确认是否已经添加过某个字母对应的结点
    var sons = [Character: TreeNode]()
    /// 从树的根到当前结点这条通路上，全部字母所组成的前缀。例如通路b->o->y，对于字母o结点而言，前缀是b；对于字母y结点而言，前缀是bo
    var prefix: String? = nil
    /// 词条的解释
    var explanation: String? = nil
    
    init(label: Character, prefix: String? = nil, explanation: String? = nil) {
        self.label = label
        self.prefix = prefix
        self.explanation = explanation
    }
}

extension TreeNode {
    func add(word: String, exp: String) {
        guard word.count > 0 else {
            if prefix != nil {
                explanation = exp
            }
            return
        }
        // 处理当前第一个字母
        let c = Character(String(word.prefix(1)))
        var found: TreeNode
        // 如果字母结点已经存在于当前父结点之下，找出它。否则就新生成一个
        if sons.keys.contains(c) {
            found = sons[c]!
        } else {
            found = TreeNode(label: c, prefix: String(prefix ?? "") + String(label), explanation: nil)
            sons[c] = found
        }
        found.add(word: String(word.suffix(word.count - 1)), exp: exp)
    }
    
    func search(word: String) -> String? {
        guard word.count > 0 else { return nil }
        var found = self
        for c in word {
            if found.sons.keys.contains(c) {
                found = found.sons[c]!
            } else {
                return nil
            }
        }
        return found.explanation
    }
}

class Trie {
    var sons = [String: Trie]()
    var isWord = false
    
    /** Inserts a word into the trie. */ //递归
    func insert0(_ word: String) {
        guard word.count > 0 else {
            isWord = true
            return
        }
        // 处理当前第一个字母
        let c = String(word.prefix(1))
        var found: Trie
        // 如果字母结点已经存在于当前父结点之下，找出它。否则就新生成一个
        if sons.keys.contains(c) {
            found = sons[c]!
        } else {
            found = Trie()
            found.isWord = false
            sons[c] = found
        }
        found.insert(String(word.suffix(word.count - 1)))
    }
    
    /** Inserts a word into the trie. */ //循环
    func insert(_ word: String) {
        var wordTemp = word
        var found: Trie = self
        for _ in 0..<word.count {
            let c = String(wordTemp.prefix(1))
            if found.sons.keys.contains(c) {
                found = found.sons[c]!
            } else {
                found.sons[c] = Trie()
                found = found.sons[c]!
            }
            wordTemp = String(wordTemp.suffix(wordTemp.count - 1))
        }
        found.isWord = true
    }
    
    /** Returns if the word is in the trie. */
    func search(_ word: String) -> Bool {
        var found = self
        for c in word {
            let key = String(c)
            if found.sons.keys.contains(key) {
                found = found.sons[key]!
            } else {
                return false
            }
        }
        return found.isWord
    }
    
    /** Returns if there is any word in the trie that starts with the given prefix. */
    func startsWith(_ prefix: String) -> Bool {
        var found = self
        for c in prefix {
            let key = String(c)
            if found.sons.keys.contains(key) {
                found = found.sons[key]!
            } else {
                return false
            }
        }
        return true
    }
}
